import sys
import numpy as np
import cv2
import math
import time
from scipy.ndimage import minimum_filter, maximum_filter, generic_filter

def boxfilter(imSrc, r):
    """优化版盒式滤波器 - 使用积分图实现"""
    hei, wid = imSrc.shape
    imDst = np.zeros_like(imSrc, dtype=np.float32)
    
    # 创建积分图
    integral = cv2.integral(imSrc.astype(np.float32))
    integral = integral[1:, 1:]  # 去掉第一行和第一列
    
    # 计算边界坐标
    r1, r2 = min(r, hei-1), min(r, wid-1)
    
    # 计算区域和
    x1 = np.arange(wid) - r2
    x2 = np.arange(wid) + r2
    y1 = np.arange(hei) - r1
    y2 = np.arange(hei) + r1
    
    # 边界处理
    x1 = np.clip(x1, 0, wid-1)
    x2 = np.clip(x2, 0, wid-1)
    y1 = np.clip(y1, 0, hei-1)
    y2 = np.clip(y2, 0, hei-1)
    
    # 使用积分图快速计算区域和
    for y in range(hei):
        for x in range(wid):
            x_min, x_max = x1[x], x2[x]
            y_min, y_max = y1[y], y2[y]
            
            # 计算区域和
            total = integral[y_max, x_max]
            if x_min > 0:
                total -= integral[y_max, x_min-1]
            if y_min > 0:
                total -= integral[y_min-1, x_max]
            if x_min > 0 and y_min > 0:
                total += integral[y_min-1, x_min-1]
            
            imDst[y, x] = total
    
    return imDst

def guidedfilter(I, p, r, eps):
    """优化版导向滤波器"""
    hei, wid = I.shape
    N = boxfilter(np.ones((hei, wid), dtype=np.float32), r)
    
    # 使用向量化计算
    mean_I = boxfilter(I, r) / N
    mean_p = boxfilter(p, r) / N
    mean_Ip = boxfilter(I * p, r) / N
    cov_Ip = mean_Ip - mean_I * mean_p
    
    mean_II = boxfilter(I * I, r) / N
    var_I = mean_II - mean_I * mean_I
    
    a = cov_Ip / (var_I + eps)
    b = mean_p - a * mean_I
    
    mean_a = boxfilter(a, r) / N
    mean_b = boxfilter(b, r) / N
    
    return mean_a * I + mean_b

def RGB2FLO_MP(img, r, g, b):
    """优化版RGB到FL颜色空间转换"""
    # 计算旋转矩阵参数
    rb_norm = math.sqrt(r*r + b*b)
    rgb_norm = math.sqrt(r*r + g*g + b*b)
    
    t = b / rb_norm if rb_norm > 0 else 0
    t2 = r / rb_norm if rb_norm > 0 else 0
    s = rb_norm / rgb_norm if rgb_norm > 0 else 0
    u = g / rgb_norm if rgb_norm > 0 else 0
    
    # 构建旋转矩阵
    Rr = np.array([[1, 0, 0], 
                  [0, s, u], 
                  [0, -u, s]])
    
    Rg = np.array([[t, 0, t2], 
                  [0, 1, 0], 
                  [-t2, 0, t]])
    
    # 应用旋转矩阵
    rgb = img.reshape(-1, 3)
    flo = rgb.dot(Rg).dot(Rr)
    flo = flo.reshape(img.shape)
    
    # 提取F和L通道
    F = flo[:, :, 2]
    L = np.sqrt(flo[:, :, 0]**2 + flo[:, :, 1]**2)
    
    return F, L, 0, flo

def FLO2RGB_MP(flo, r, g, b):
    """优化版FL到RGB颜色空间转换"""
    # 计算旋转矩阵参数
    rb_norm = math.sqrt(r*r + b*b)
    rgb_norm = math.sqrt(r*r + g*g + b*b)
    
    t = b / rb_norm if rb_norm > 0 else 0
    t2 = r / rb_norm if rb_norm > 0 else 0
    s = rb_norm / rgb_norm if rgb_norm > 0 else 0
    u = g / rgb_norm if rgb_norm > 0 else 0
    
    # 构建旋转矩阵
    Rr = np.array([[1, 0, 0], 
                  [0, s, u], 
                  [0, -u, s]])
    
    Rg = np.array([[t, 0, t2], 
                  [0, 1, 0], 
                  [-t2, 0, t]])
    
    # 应用逆旋转
    flt = flo.reshape(-1, 3)
    rgb = flt.dot(Rr.T).dot(Rg.T)
    return rgb.reshape(flo.shape)

def find_A(F, L, img, scale=160):
    """优化版大气光估计 - 使用滑动窗口滤波"""
    # 下采样图像
    h, w = F.shape
    aspect = h / w
    new_h = scale
    new_w = int(new_h / aspect)
    
    F_resized = cv2.resize(F, (new_w, new_h), interpolation=cv2.INTER_AREA)
    L_resized = cv2.resize(L, (new_w, new_h), interpolation=cv2.INTER_AREA)
    img_resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
    
    # 使用滑动窗口滤波
    A_wind = 18
    
    # 使用高效的最小/最大滤波器
    F_Amin = minimum_filter(F_resized, size=(A_wind, A_wind), mode='reflect')
    L_Amax = maximum_filter(L_resized, size=(A_wind, A_wind), mode='reflect')
    
    # 计算候选点
    pix_cnt = F_Amin.size
    light_cnt = max(1, int(pix_cnt * 0.001))  # 确保至少1个点
    
    # 获取F_Amin中最大的0.1%的点
    threshold = np.partition(F_Amin.flatten(), -light_cnt)[-light_cnt]
    F_mask = (F_Amin >= threshold)
    
    # 获取L_Amax中最大的10%的点
    L_vals = L_Amax[F_mask]
    if L_vals.size > 0:
        L_cnt = max(1, int(len(L_vals) * 0.1))
        L_threshold = np.partition(L_vals, -L_cnt)[-L_cnt]
        candidate_mask = F_mask & (L_Amax >= L_threshold)
    else:
        candidate_mask = np.zeros_like(F_mask, dtype=bool)
    
    # 获取候选点坐标
    y_coords, x_coords = np.where(candidate_mask)
    
    # 计算天空颜色和大气光值
    if len(y_coords) > 0:
        R_sky = np.mean(img_resized[y_coords, x_coords, 2])
        G_sky = np.mean(img_resized[y_coords, x_coords, 1])
        B_sky = np.mean(img_resized[y_coords, x_coords, 0])
        A = np.mean(F_resized[y_coords, x_coords])
    else:
        # 如果没有候选点，使用图像最亮的1%像素
        flat_F = F_resized.flatten()
        bright_threshold = np.percentile(flat_F, 99)
        bright_pixels = F_resized >= bright_threshold
        
        R_sky = np.mean(img_resized[bright_pixels, 2])
        G_sky = np.mean(img_resized[bright_pixels, 1])
        B_sky = np.mean(img_resized[bright_pixels, 0])
        A = np.mean(F_resized[bright_pixels])
    
    return R_sky, G_sky, B_sky, A

def FL_dehaze(flo, t, A):
    """优化版FL空间去雾 - 使用向量化操作"""
    # 裁剪透射率到合理范围
    t_clipped = np.clip(t, 0.01, 1.0)
    
    # 创建输出数组
    dehaze_flo = np.empty_like(flo)
    
    # 向量化操作
    dehaze_flo[:, :, 0] = flo[:, :, 0] / t_clipped[:, :]  # L通道1
    dehaze_flo[:, :, 1] = flo[:, :, 1] / t_clipped[:, :]  # L通道2
    dehaze_flo[:, :, 2] = (flo[:, :, 2] - A) / t_clipped[:, :] + A  # F通道
    
    return dehaze_flo

def redefine_d(F, d):
    """优化版距离图重定义"""
    r = 15
    eps = 0.1
    
    # 使用最小值滤波
    d_wind = minimum_filter(d, size=2, mode='reflect')
    
    # 使用导向滤波
    return guidedfilter(F, d_wind, r, eps)

def dehaze_main(image_path, output_path=None, show_result=True):
    """优化版主函数 - 图像去雾流程"""
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图像: {image_path}")
    
    # 转换为浮点数并归一化
    hazeimg = img.astype(np.float32) / 255.0
    height, width = hazeimg.shape[:2]
    
    start_time = time.time()
    
    # 第一次大气光估计
    print("第一次大气光估计...")
    F1, L1, _, _ = RGB2FLO_MP(hazeimg, 1, 1, 1)
    R_sky, G_sky, B_sky, _ = find_A(F1, L1, hazeimg)
    
    # 第二次大气光估计
    print("第二次大气光估计...")
    F2, L2, _, _ = RGB2FLO_MP(hazeimg, R_sky, G_sky, B_sky)
    R_sky1, G_sky1, B_sky1, A = find_A(F2, L2, hazeimg)
    
    # 最终颜色空间转换
    print("颜色空间转换...")
    F, L, _, flo = RGB2FLO_MP(hazeimg, R_sky1, G_sky1, B_sky1)
    
    print(f"预处理时间: {time.time()-start_time:.2f}秒")
    
    # 计算深度图和透射率
    print("计算深度图和透射率...")
    norm_factor = math.sqrt(R_sky1**2 + G_sky1**2 + B_sky1**2)
    F_one = F / norm_factor if norm_factor > 0 else F
    L_one = L / np.max(L) if np.max(L) > 0 else L
    
    # 使用改进的深度估计公式
    d = 1.016 * F_one - 0.589 * L_one + 0.12413
    
    # 优化透射率计算
    t = np.exp(-0.75 * d)
    t = np.clip(t, 0.1, 0.99)  # 更合理的裁剪范围
    
    print(f"透射率计算时间: {time.time()-start_time:.2f}秒")
    
    # FL空间去雾
    print("FL空间去雾...")
    dehaze_flo = FL_dehaze(flo, t, A)
    
    # 转换回RGB空间
    print("转换回RGB空间...")
    rgb_d = FLO2RGB_MP(dehaze_flo, R_sky1, G_sky1, B_sky1)
    rgb_d = np.clip(rgb_d, 0, 1)
    
    # 后处理 - 对比度增强
    print("后处理...")
    dehazed = (rgb_d * 255).astype(np.uint8)
    dehazed = cv2.detailEnhance(dehazed, sigma_s=10, sigma_r=0.15)
    
    total_time = time.time() - start_time
    print(f"总处理时间: {total_time:.2f}秒")
    
    # 显示结果
    # if show_result:
    #     cv2.imshow("Original", img)
    #     cv2.imshow("Dehazed", dehazed)
    #     cv2.waitKey(0)
    #     cv2.destroyAllWindows()
    
    # 保存结果
    if output_path:
        cv2.imwrite(output_path, dehazed)
        print(f"结果已保存至: {output_path}")
    
    return dehazed

def batch_process(images, output_dir):
    """批量处理图像"""
    from pathlib import Path
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True)
    
    results = []
    for img_path in images:
        output_path = output_dir / f"dehazed_{Path(img_path).name}"
        try:
            result = dehaze_main(str(img_path), str(output_path), show_result=False)
            results.append(result)
        except Exception as e:
            print(f"处理 {img_path} 时出错: {e}")
    
    return results

if __name__ == "__main__":
    import argparse
    import os
    
    # parser = argparse.ArgumentParser(description="图像去雾算法")
    # parser.add_argument("input", help="输入图像路径或目录")
    # parser.add_argument("-o", "--output", help="输出目录")
    # parser.add_argument("-b", "--batch", action="store_true", help="批量处理模式")
    
    # args = parser.parse_args()
    
    # if args.batch:
    #     # 批量处理模式
    #     if os.path.isdir(args.input):
    #         images = [os.path.join(args.input, f) for f in os.listdir(args.input) 
    #                  if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    #         output_dir = args.output or os.path.join(args.input, "dehazed_results")
    #         batch_process(images, output_dir)
    #     else:
    #         print("错误: 批量处理需要输入目录")
    # else:
    #     # 单图像处理模式
        # output_path = args.output or f"dehazed_{os.path.basename(args.input)}"
        # dehaze_main(args.input, output_path)
    if len(sys.argv) != 3:
        print("Usage: python SingMain.py <input_path> <output_path>")
        sys.exit(1)
    input_path = sys.argv[1]
    output_path = sys.argv[2]
    dehaze_main(input_path, output_path)